package tree

import (
	"fmt"
	"gitee.com/kklt1996/data-structure/queue"
	"gitee.com/kklt1996/data-structure/stack"
	"gitee.com/kklt1996/data-structure/util"
	"strings"
)

/*
	伸展树:伸展树也是二分搜索树
	二分搜索树是二叉树,因为其中序列遍历的有序性也称二分搜索树为排序树

	二分搜索树的每隔节点的值:
		大于其左子节点的所有节点的值
		小于其右子节点的所有节点的值
	为了满足此特性,二分搜索树中不允许有相等的元素

	每一棵子树也是二分搜索树

	存入二分搜索树中的节点必须是可以比较的,想要存入基础类型就必须给基础类型实现
	 common.CompareAble接口
*/
type SplayNode struct {
	key, value  interface{}
	left, right *SplayNode
}

/*
	伸展树: 伸展树的特性是将最近集中访问的元素移动到树的头部.虽然curd的均摊时间复杂度是O(n)的.
*/
type SplayTree struct {
	// 伸展树的根节点
	root *SplayNode
	// 元素元素个数
	size int

	/*
		使用comparator比较堆中元素大小,实现common.CompareAble接口和传入比较器形式二选一
		1 thisKey＞compareKey
		0 thisKey=compareKey
		-1 thisKey<compareKey
	*/
	comparator func(thisKey interface{}, compareKey interface{}) int
}

/*
	O(1)
	获取元素个数
*/
func (tree SplayTree) GetSize() int {
	return tree.size
}

/*
	O(1)
	判断是否为空
*/
func (tree SplayTree) IsEmpty() bool {
	return tree.size == 0
}

/*
	O(h)
	添加元素
*/
func (tree *SplayTree) Add(key interface{}, value interface{}) {
	tree.root, _, _ = tree.add(tree.root, key, value)
}

/*
	向以node为根节点的二分搜索树中插入元素e
	返回以node为根节点的二分搜索树插入新节点之后以的根
	空节点也是一个二叉树
	依次返回newRoot, addNode, needSplayFlag
*/
func (tree *SplayTree) add(root *SplayNode, key interface{}, value interface{}) (*SplayNode, *SplayNode, bool) {
	// 找到满足条件,且为空的位置,就插入元素返回根节点,也就是根节点赋值为新增节点的地址
	if root == nil {
		tree.size++
		// 这就是递归的临界点,在最小的树中找到了插入元素的位置
		addNode := &SplayNode{key: key, value: value}
		return addNode, addNode, false
	} else {
		var addNode *SplayNode
		var needSplayFlag bool
		// 要添加的元素和当前根节点比较大小
		i := tree.compare(key, root.key)
		if i < 0 {
			root.left, addNode, needSplayFlag = tree.add(root.left, key, value)
			return tree.splay(root, addNode, needSplayFlag), addNode, !needSplayFlag
		} else if i > 0 {
			root.right, addNode, needSplayFlag = tree.add(root.right, key, value)
			return tree.splay(root, addNode, needSplayFlag), addNode, !needSplayFlag
		} else {
			// 相等
			root.value = value
			return root, root, false
		}
	}
}

/*
	O(h)
	查询二分搜索树中是否包含某个元素
*/
func (tree *SplayTree) Contains(key interface{}) bool {
	var getNode *SplayNode
	tree.root, getNode, _ = tree.get(tree.root, key)
	if getNode == nil {
		return false
	} else {
		return true
	}
}

/*
	查询key对应的value,在get的过程中重新连接节点,以便加入二层伸展操作
*/
func (tree *SplayTree) Get(key interface{}) interface{} {
	var getNode *SplayNode
	tree.root, getNode, _ = tree.get(tree.root, key)

	if getNode == nil {
		return nil
	} else {
		return getNode.value
	}
}

/*
	返回新的节点和要查找的节点
*/
func (tree *SplayTree) get(root *SplayNode, key interface{}) (newRoot *SplayNode, getNode *SplayNode, needSplay bool) {
	if root == nil {
		return nil, nil, false
	}
	i := tree.compare(key, root.key)
	if i == 0 {
		// 如果本身就是根节点，则肯定不需要伸展操作
		return root, root, false
	} else if i < 0 {
		root.left, getNode, needSplay = tree.get(root.left, key)
	} else {
		root.right, getNode, needSplay = tree.get(root.right, key)
	}
	// 对节进行双层伸展
	return tree.splay(root, getNode, needSplay), getNode, !needSplay
}

/*
	前序遍历,进行一些操作
*/
func (tree SplayTree) PreOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	tree.preOrder(tree.root, operatorFunc)
}

/*
	以某个节点为根节点进行前序列遍历
*/
func (tree SplayTree) preOrder(root *SplayNode, operatorFunc func(key interface{}, value interface{}) bool) {
	// 当没有node的时候表示结束
	if root == nil {
		return
	}
	if operatorFunc(root.key, root.value) {
		return
	}
	tree.preOrder(root.left, operatorFunc)
	tree.preOrder(root.right, operatorFunc)
}

/*
	二分搜索树的中序遍历,中序遍历的结果是升序排序的
*/
func (tree SplayTree) InOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	tree.inOrder(tree.root, operatorFunc)
}

/*
	以某个node为根节点进行中序列遍历
*/
func (tree SplayTree) inOrder(root *SplayNode, operatorFunc func(key interface{}, value interface{}) bool) {
	// 当没有node的时候表示结束
	if root == nil {
		return
	}
	tree.inOrder(root.left, operatorFunc)
	if operatorFunc(root.key, root.value) {
		return
	}
	tree.inOrder(root.right, operatorFunc)
}

/*
	二分搜索树的后序遍历
*/
func (tree SplayTree) PostOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	tree.postOrder(tree.root, operatorFunc)
}

/*
	以某个node为根节点进行后序列遍历
*/
func (tree SplayTree) postOrder(root *SplayNode, operatorFunc func(key interface{}, value interface{}) bool) {
	// 当没有node的时候表示结束
	if root == nil {
		return
	}
	tree.postOrder(root.left, operatorFunc)
	tree.postOrder(root.right, operatorFunc)
	if operatorFunc(root.key, root.value) {
		return
	}
}

/*
	二分搜索树的前序遍历非递归实现
*/
func (tree SplayTree) PreOrderNR(operatorFunc func(key interface{}, value interface{}) bool) {
	var stk stack.Stack = stack.CreateLinkedListStack()
	if tree.root != nil {
		stk.Push(tree.root)
	}

	for !stk.IsEmpty() {
		pop, _ := stk.Pop()
		node := pop.(*SplayNode)
		if !operatorFunc(node.key, node.value) {
			return
		}
		if node.right != nil {
			stk.Push(node.right)
		}
		if node.left != nil {
			stk.Push(node.left)
		}
	}

}

/*
	借助队列实现二分搜索树的层序遍历
*/
func (tree SplayTree) LevelOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	var q queue.Queue = queue.CreateLinkedListQueue()
	if tree.root != nil {
		q.Enqueue(tree.root)
	}
	for !q.IsEmpty() {
		// 遍历当前层节点
		dequeue, _ := q.Dequeue()
		node := dequeue.(*SplayNode)
		// 搜索结束,直接返回结果
		if operatorFunc(node.key, node.value) {
			return
		}

		// 下一层节点入队,当前层节点遍历完成后就会遍历下一层节点
		if node.left != nil {
			q.Enqueue(node.left)
		}
		if node.right != nil {
			q.Enqueue(node.right)
		}
	}
}

/*
	O(h)
	查找二分搜索树中最小的元素
*/
func (tree SplayTree) Minimum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, BstIsEmptyError{}
	}
	minimumNode := tree.minimum(tree.root)
	return minimumNode.key, minimumNode.value, nil
}

/*
	递归查找最小元素的节点,也就是微分搜索树最左边的节点
*/
func (tree SplayTree) minimum(node *SplayNode) *SplayNode {
	if node.left == nil {
		return node
	}
	return tree.minimum(node.left)
}

/*
	O(h)
	查找最小的元素
*/
func (tree SplayTree) Maximum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, BstIsEmptyError{}
	}
	maximumNode := tree.maximum(tree.root)
	return maximumNode.key, maximumNode.value, nil
}

/*
	递归查找最大元素的节点,也就是微分搜索树最右边的节点
*/
func (tree SplayTree) maximum(node *SplayNode) *SplayNode {
	if node.right == nil {
		return node
	}
	return tree.maximum(node.right)
}

/*
	O(h)
	删除二分搜索树的最小的节点
*/
func (tree *SplayTree) RemoveMinimum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, BstIsEmptyError{}
	}
	newRoot, minimumKey, minimumValue := tree.removeMinimum(tree.root)
	// 更新以 tree.root 为根节点子树的根为新的根
	tree.root = newRoot
	return minimumKey, minimumValue, nil
}

/*
	删除以node为根节点的二分搜索树的最小值
	返回删除最小的元素节点后二分搜索树的根和最小的元素
*/
func (tree *SplayTree) removeMinimum(root *SplayNode) (*SplayNode, interface{}, interface{}) {
	if root.left == nil {
		// 找到最小的元素,获取要返回的右子树
		rightTree := root.right
		// 删除节点和右孩子的关系
		root.right = nil
		// 元素个数-1
		tree.size--
		// 返回最小节点的右子树,和最小节点的值
		return rightTree, root.key, root.value
	} else {
		// 删除左子树最小的元素
		newRoot, minimumKey, minimumValue := tree.removeMinimum(root.left)
		// 更新以 root.left 为根节点子树的根为新的根
		root.left = newRoot
		// 返回原根和删除最小节点的value
		return root, minimumKey, minimumValue
	}
}

/*
	O(h)
	删除二分搜索树的最大的节点
*/
func (tree *SplayTree) RemoveMaximum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, BstIsEmptyError{}
	}
	newRoot, maximumKey, maximumValue := tree.removeMaximum(tree.root)
	tree.root = newRoot
	return maximumKey, maximumValue, nil
}

/*
	删除以node为根节点的二分搜索树的最大值的节点
	返回删除最大的元素节点后二分搜索树的根
*/
func (tree *SplayTree) removeMaximum(root *SplayNode) (*SplayNode, interface{}, interface{}) {
	if root.right == nil {
		// 找到最大的元素,获取要返回的左子树
		leftTree := root.left
		// 删除节点和左孩子的关系
		root.left = nil
		// 元素个数-1
		tree.size--
		// 返回最大节点的左子树
		return leftTree, root.key, root.value
	} else {
		// 删除右子树最大的元素
		newRoot, maximumKey, maximumValue := tree.removeMaximum(root.right)
		root.right = newRoot
		return root, maximumKey, maximumValue
	}
}

/*
	O(h)
	删除二分搜索树中任意节点,删除节点的时候进行伸展操作主要是为了压缩树的高度
*/
func (tree *SplayTree) RemoveElement(key interface{}) (interface{}, error) {
	if tree.IsEmpty() {
		return nil, SplayIsEmptyError{}
	}
	var getNode *SplayNode
	// 如果找到元素将元素移动到头部
	tree.root, getNode, _ = tree.get(tree.root, key)
	if getNode == nil {
		return nil, nil
	} else {
		// 将移动到头部的元素进行删除
		tree.root, _ = tree.removeElement(tree.root, key)
		return getNode.value, nil
	}
}

/*
	删除以某个node为根节点的二分搜索树中的某一个节点
	返回删除指定节点后新二分搜索树的根
*/
func (tree *SplayTree) removeElement(root *SplayNode, key interface{}) (*SplayNode, interface{}) {
	// 根节点为空直接返回
	if root == nil {
		return nil, nil
	}
	i := tree.compare(key, root.key)
	if i == 0 {
		// 递归终止条件
		// 找到需要删除的节点
		if root.left == nil {
			// 左子树为空,那么直接将右节点升级为二分搜索树的根
			rightTree := root.right
			root.right = nil
			tree.size--
			return rightTree, root.value
		}
		if root.right == nil {
			// 右子树为空,那么直接将左子树升级为二分搜索树的根
			leftTree := root.left
			root.left = nil
			tree.size--
			return leftTree, root.value
		}
		// 左右子树都不为空,那么为满足二分搜索树的性质
		// 需要从右子树中删除最小的节点或者从左子树中删除最大的节点
		// 将删除的节点赋值给当前根节点
		rightNewRoot, removeKey, removeValue := tree.removeMinimum(root.right)
		// 接收删除的值
		retRemoveValue := root.value
		root.right = rightNewRoot
		root.key = removeKey
		root.value = removeValue
		return root, retRemoveValue
	} else {
		// 用更小的结果拼凑最终的结果,在子树中删除元素
		var retRemoveValue interface{}
		if i < 0 {
			root.left, retRemoveValue = tree.removeElement(root.left, key)
		} else {
			root.right, retRemoveValue = tree.removeElement(root.right, key)
		}
		// 根不变
		return root, retRemoveValue
	}
}

/*
	进行双层伸展操作,返回新的根节点
*/
func (tree SplayTree) splay(root *SplayNode, splayNode *SplayNode, needSplay bool) *SplayNode {
	if splayNode == nil {
		// 不需要伸展操作,返回原来的根节点
		return root
	}
	if needSplay {
		// 需要伸展的时候一定是具有祖先节点
		if root.left != nil && root.left.left == splayNode {
			// LL
			// 对祖先节点进行右旋转
			splayParentNode := tree.rightRotate(root)
			// 对父亲节点进行右旋转
			return tree.rightRotate(splayParentNode)
		} else if root.left != nil && root.left.right == splayNode {
			// LR
			// 对父亲节点进行左旋转
			root.left = tree.leftRotate(root.left)
			// 对祖先节点进行右旋转
			return tree.rightRotate(root)
		} else if root.right != nil && root.right.right == splayNode {
			// RR
			// 对祖先节点进行左旋转
			splayParentNode := tree.leftRotate(root)
			// 对父亲节点进行左旋转
			return tree.leftRotate(splayParentNode)
		} else {
			// RL
			// 对父亲节点进行左旋转
			root.right = tree.rightRotate(root.right)
			// 对祖先节点进行右旋转
			return tree.leftRotate(root)
		}
	} else {
		if root == tree.root && root.left == splayNode {
			// 标识不需要伸展的时候可能是这种情况.父亲节点是根节点,也需要进行伸展操作。
			// 右旋转
			return tree.rightRotate(root)
		} else if root == tree.root && root.right == splayNode {
			// 标识不需要伸展的时候可能是这种情况.父亲节点是根节点,也需要进行伸展操作。
			// 左旋转
			return tree.leftRotate(root)
		} else {
			// 真的不需要伸展
			return root
		}
	}
}

/*
		右旋转操作
		 y				 x
		/ \				/ \
	   x   T4------>   z   y
      / \             / \  / \
	 z   T3          T1 T2 T3 T4
    / \
   T1 T2
	x,y,z 节点都可能有左右子树
*/
func (tree SplayTree) rightRotate(root *SplayNode) *SplayNode {
	// 获取要返回的根节点
	newRoot := root.left
	root.left.right, root.left = root, root.left.right
	// 回溯更新左右孩子节点发生变化的节点
	return newRoot
}

/*
	左旋转操作
			y					x
			 \				   / \
			  x     ----->    y   z
               \
                z
*/
func (tree SplayTree) leftRotate(root *SplayNode) *SplayNode {
	newRoot := root.right
	root.right.left, root.right = root, root.right.left
	return newRoot
}

/*
	获取树的最大高度
*/
func (tree SplayTree) MaxDepth() int {
	return tree.maxDepth(tree.root)
}

/*
	获取二分搜索树的最大高度
	返回当前二分搜索树的高度
*/
func (tree SplayTree) maxDepth(root *SplayNode) int {
	if root == nil {
		// 空节点返回0,不计算高度
		return 0
	} else {
		leftTreeMaxDepth := tree.maxDepth(root.left)
		rightTreeMaxDepth := tree.maxDepth(root.right)
		// 左右子树的最大高度+1,就是当前二分搜索树的最大高度
		// +1 表示每经过一个非空节点，高度+1
		if leftTreeMaxDepth > rightTreeMaxDepth {
			return leftTreeMaxDepth + 1
		} else {
			return rightTreeMaxDepth + 1
		}
	}
}

func (tree SplayTree) compare(thisValue interface{}, compareValue interface{}) int {
	if tree.comparator != nil {
		return tree.comparator(thisValue, compareValue)
	} else {
		return util.DefaultComparator(thisValue, compareValue)
	}
}

func (tree SplayTree) String() string {
	res := "["
	tree.inOrder(tree.root, func(key interface{}, value interface{}) bool {
		res += fmt.Sprint(key) + ","
		return false
	})
	res = strings.TrimSuffix(res, ",")
	res += "]"
	return res
}

type SplayIsEmptyError struct {
}

func (SplayIsEmptyError) Error() string {
	return "splayTree is empty"
}
